# -*- coding: utf-8 -*-
import numpy as np
import json
import const
from scipy import sparse
from utils.utils import npArrayToCPPFile

class Reservoir(object):
	def __init__(self, loadFromConst=False, modelDir=const.PATH_TO_DIR):
		if loadFromConst:
			if modelDir != "''":
				load_file = modelDir + 'echo.json'
			else:
				load_file = const.PATH_TO_DIR + 'echo.json'
			loaded_data = {}
			with open(load_file, 'r') as file:
				loaded_data = json.load(file)
			self._n_internal_units = loaded_data['_n_internal_units']
			self._input_scaling = loaded_data['_input_scaling']
			self._noise_level = loaded_data['_noise_level']
			self._leak = loaded_data['_leak']
			self._input_weights = np.array(loaded_data['_input_weights'])
			self._internal_weights = np.array(loaded_data['_internal_weights'])
			print ('echo load from %s'%load_file)
		else:
			self._n_internal_units = const.N_INTERNAL_UNITS
			self._input_scaling = const.INPUT_SCALING
			self._noise_level = const.NOISE_LEVEL
			self._leak = const.LEAK
			self._input_weights = None
			if const.CIRCLE:
				self._internal_weights = self._initialize_internal_weights_Circ(const.N_INTERNAL_UNITS, const.SPECTRAL_RADIUS)
			else:
				self._internal_weights = self._initialize_internal_weights(const.N_INTERNAL_UNITS, const.CONNECTIVITY, const.SPECTRAL_RADIUS)


	def _initialize_internal_weights_Circ(self, n_internal_units, spectral_radius):
		internal_weights = np.zeros((n_internal_units, n_internal_units))
		internal_weights[0,-1] = 1.0
		for i in range(n_internal_units-1):
			internal_weights[i+1,i] = 1.0
		E, _ = np.linalg.eig(internal_weights)
		e_max = np.max(np.abs(E))
		internal_weights /= np.abs(e_max)/spectral_radius 
				
		return internal_weights
	
	
	def _initialize_internal_weights(self, n_internal_units, connectivity, spectral_radius):
		internal_weights = sparse.rand(n_internal_units, n_internal_units, density=connectivity).todense()
		internal_weights[np.where(internal_weights > 0)] -= 0.5
		E, _ = np.linalg.eig(internal_weights)
		e_max = np.max(np.abs(E))
		internal_weights /= np.abs(e_max)/spectral_radius	   
		return internal_weights


	def _compute_state_matrix(self, X, n_drop=0):
		N, T, _ = X.shape
		previous_state = np.zeros((N, self._n_internal_units), dtype=float)
		state_matrix = np.empty((N, T - n_drop, self._n_internal_units), dtype=float)
		for t in range(T):
			current_input = X[:, t, :]
			state_before_tanh = self._internal_weights.dot(previous_state.T) + self._input_weights.dot(current_input.T)
			state_before_tanh += np.random.rand(self._n_internal_units, N)*self._noise_level
			if self._leak is None:
				previous_state = np.tanh(state_before_tanh).T
			else:
				previous_state = (1.0 - self._leak)*previous_state + np.tanh(state_before_tanh).T
			if (t > n_drop - 1):
				state_matrix[:, t - n_drop, :] = previous_state
		return state_matrix


	def getStates(self, X, n_drop=0, bidir=True):
		N, T, V = X.shape
		if self._input_weights is None:
			self._input_weights = (2.0*np.random.binomial(1, 0.5 , [self._n_internal_units, V]) - 1.0)*self._input_scaling
		states = self._compute_state_matrix(X, n_drop)
		if bidir is True:
			X_r = X[:, ::-1, :]
			states_r = self._compute_state_matrix(X_r, n_drop)
			states = np.concatenate((states, states_r), axis=2)
		return states
	
	def saveModel(self, path):
		param = {
			'_input_weights': self._input_weights.tolist(),
			'_internal_weights': self._internal_weights.tolist(),
			'_n_internal_units': self._n_internal_units,
			'_input_scaling': self._input_scaling,
			'_noise_level': self._noise_level,
			'_leak': self._leak,
		}
		save_path = path+'echo.json'
		with open(save_path, 'w') as file:
			json.dump(param, file)
	
	def saveParamsToCppFile(self):
		npArrayToCPPFile(self._input_weights, '_input_weights', const.CPP_CONST_DIR+'echo_input_weights.h')
		npArrayToCPPFile(self._internal_weights, '_internal_weights', const.CPP_CONST_DIR+'echo_internal_weights.h')
		npArrayToCPPFile(np.array([[self._n_internal_units, self._input_scaling, self._noise_level, self._leak]]), '_all_params', const.CPP_CONST_DIR+'echo_all_params.h')